#!/usr/bin/env python3
"""
Dolphin文档解析API客户端示例
演示如何调用API进行文档解析
"""

import base64
import json
import requests
import time
from pathlib import Path


class DolphinAPIClient:
    def __init__(self, base_url: str = "http://localhost:8000"):
        """
        初始化API客户端
        
        Args:
            base_url: API服务器地址
        """
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
    
    def health_check(self) -> dict:
        """检查API服务健康状态"""
        try:
            response = self.session.get(f"{self.base_url}/health")
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            return {"error": str(e)}
    
    def parse_file(self, file_path: str, max_batch_size: int = 4) -> dict:
        """
        上传并解析文档文件
        
        Args:
            file_path: 文档文件路径
            max_batch_size: 批处理大小
            
        Returns:
            解析结果字典
        """
        file_path = Path(file_path)
        
        if not file_path.exists():
            return {"error": f"文件不存在: {file_path}"}
        
        try:
            with open(file_path, 'rb') as f:
                files = {'file': (file_path.name, f, self._get_content_type(file_path))}
                data = {'max_batch_size': max_batch_size}
                
                response = self.session.post(
                    f"{self.base_url}/parse",
                    files=files,
                    data=data
                )
                response.raise_for_status()
                return response.json()
        
        except requests.RequestException as e:
            return {"error": str(e)}
    
    def parse_base64(self, image_path: str, max_batch_size: int = 4) -> dict:
        """
        通过Base64编码解析图像
        
        Args:
            image_path: 图像文件路径
            max_batch_size: 批处理大小
            
        Returns:
            解析结果字典
        """
        image_path = Path(image_path)
        
        if not image_path.exists():
            return {"error": f"文件不存在: {image_path}"}
        
        # 验证是否为图像文件
        if image_path.suffix.lower() not in ['.jpg', '.jpeg', '.png']:
            return {"error": "Base64模式仅支持JPG、JPEG、PNG格式"}
        
        try:
            # 读取并编码图像
            with open(image_path, 'rb') as f:
                image_data = base64.b64encode(f.read()).decode('utf-8')
            
            data = {
                'image_data': image_data,
                'filename': image_path.name,
                'max_batch_size': max_batch_size
            }
            
            response = self.session.post(
                f"{self.base_url}/parse_base64",
                data=data
            )
            response.raise_for_status()
            return response.json()
        
        except requests.RequestException as e:
            return {"error": str(e)}
    
    def _get_content_type(self, file_path: Path) -> str:
        """根据文件扩展名获取Content-Type"""
        ext = file_path.suffix.lower()
        content_types = {
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.png': 'image/png',
            '.pdf': 'application/pdf'
        }
        return content_types.get(ext, 'application/octet-stream')


def print_parse_result(result: dict):
    """打印解析结果的格式化输出"""
    if "error" in result:
        print(f"❌ 错误: {result['error']}")
        return
    
    if not result.get("success", False):
        print(f"❌ 解析失败: {result.get('message', '未知错误')}")
        return
    
    print(f"✅ 解析成功!")
    print(f"📄 任务ID: {result['task_id']}")
    print(f"📑 总页数: {result['total_pages']}")
    print(f"⏱️  处理时间: {result['processing_time']}秒")
    
    # 打印每页的解析结果
    for page in result['results']:
        print(f"\n📖 第{page['page_number']}页 - 共{len(page['elements'])}个元素:")
        
        for i, element in enumerate(page['elements'], 1):
            element_type = {
                'para': '📝 段落',
                'tab': '📊 表格',
                'fig': '🖼️  图片'
            }.get(element['label'], f"❓ {element['label']}")
            
            print(f"  {i:2d}. {element_type} (顺序: {element['reading_order']})")
            
            # 打印内容预览（限制长度）
            content = element['text']
            if len(content) > 100:
                content = content[:100] + "..."
            
            print(f"      内容: {content}")
            
            if element.get('figure_path'):
                print(f"      图片路径: {element['figure_path']}")


def main():
    """主函数 - 演示API使用"""
    print("🐬 Dolphin文档解析API客户端示例")
    print("=" * 50)
    
    # 初始化客户端
    client = DolphinAPIClient("http://localhost:8000")
    
    # 检查服务健康状态
    print("🔍 检查API服务状态...")
    health = client.health_check()
    
    if "error" in health:
        print(f"❌ 无法连接到API服务: {health['error']}")
        print("请确保API服务正在运行 (python api_server.py)")
        return
    
    print(f"✅ API服务状态: {health['status']}")
    print(f"🤖 模型加载状态: {'已加载' if health['model_loaded'] else '未加载'}")
    
    if not health['model_loaded']:
        print("❌ 模型未加载，请检查配置和模型文件")
        return
    
    # 示例1: 解析图像文件
    print("\n" + "=" * 50)
    print("📸 示例1: 解析图像文件")
    
    # 你可以替换为实际的图像文件路径
    image_file = "./demo/page_imgs/page_1.jpeg"
    
    if Path(image_file).exists():
        print(f"正在解析图像: {image_file}")
        start_time = time.time()
        result = client.parse_file(image_file, max_batch_size=4)
        end_time = time.time()
        
        print(f"客户端总耗时: {end_time - start_time:.2f}秒")
        print_parse_result(result)
    else:
        print(f"⚠️  示例图像文件不存在: {image_file}")
    
    # 示例2: Base64方式解析
    print("\n" + "=" * 50)
    print("🔤 示例2: Base64方式解析")
    
    if Path(image_file).exists():
        print(f"正在通过Base64解析图像: {image_file}")
        start_time = time.time()
        result = client.parse_base64(image_file, max_batch_size=4)
        end_time = time.time()
        
        print(f"客户端总耗时: {end_time - start_time:.2f}秒")
        print_parse_result(result)
    else:
        print(f"⚠️  示例图像文件不存在: {image_file}")
    
    # 示例3: 解析PDF文件
    print("\n" + "=" * 50)
    print("📄 示例3: 解析PDF文件")
    
    pdf_file = "./demo/page_imgs/page_6.pdf"
    
    if Path(pdf_file).exists():
        print(f"正在解析PDF: {pdf_file}")
        start_time = time.time()
        result = client.parse_file(pdf_file, max_batch_size=4)
        end_time = time.time()
        
        print(f"客户端总耗时: {end_time - start_time:.2f}秒")
        print_parse_result(result)
    else:
        print(f"⚠️  示例PDF文件不存在: {pdf_file}")
    
    print("\n🎉 API客户端示例完成!")


if __name__ == "__main__":
    main()